#!/usr/bin/env python3
"""
Author: Logan.Li
Gitee: https://gitee.com/attacker
email: admin@attacker.club
Date: 2025-01-05 12:25:52
LastEditTime: 2025-09-17
Description: Kubernetes资源清单统计工具 - 收集并分析集群中Deployment的资源配置信息，生成详细的Excel报告
"""

import os
import sys
import subprocess
from collections import defaultdict

# 检查并安装必要的模块
def check_and_install_dependencies():
    """检查并安装必要的Python模块"""
    required_modules = {
        'kubernetes': 'kubernetes',
        'openpyxl': 'openpyxl',
        'urllib3': 'urllib3'
    }
    
    missing_modules = []
    for module_name, package_name in required_modules.items():
        try:
            __import__(module_name)
        except ImportError:
            missing_modules.append(package_name)
    
    if missing_modules:
        print(f"缺少必要模块: {', '.join(missing_modules)}")
        print("正在尝试安装所需模块...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install"] + missing_modules)
            print("模块安装完成，请重新运行脚本")
        except subprocess.CalledProcessError as e:
            print(f"模块安装失败: {e}")
            print("请手动安装: pip install " + " ".join(missing_modules))
            sys.exit(1)
        sys.exit(0)

# 执行依赖检查
check_and_install_dependencies()

# 导入所需模块
from kubernetes import client, config
from openpyxl import Workbook
import urllib3

# 禁用 SSL 警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# 需要放在底部的系统命名空间
SYSTEM_NAMESPACES = {'ingress-nginx', 'kube-system', 'kuboard', 'prometheus','devops','monitor','logging','istio-system','istio-egressgateway','istio-ingressgateway'}

def setup_kubernetes_client():
    """设置 Kubernetes 客户端连接"""
    try:
        # 尝试加载 kubeconfig 配置
        config.load_kube_config()
        api_client = client.ApiClient()
        return client.CoreV1Api(api_client), client.AppsV1Api(api_client)
    except Exception as e:
        print(f"加载 kubeconfig 失败: {e}")
        sys.exit(1)

def extract_deployment_name(pod_name):
    """从 Pod 名称中提取 Deployment 名称（优化版）"""
    if not pod_name or '-' not in pod_name:
        return pod_name
    
    parts = pod_name.split('-')
    
    # 匹配标准命名模式：前缀-哈希-随机后缀
    if len(parts) >= 3:
        # 检查最后两部分是否符合哈希和随机后缀的特征
        last_part = parts[-1]
        second_last = parts[-2]
        
        # 随机后缀通常是5个字符（字母数字混合）
        is_random_suffix = (len(last_part) == 5 and last_part.isalnum())
        
        # 模板哈希通常是9-10个字符（数字和小写字母）
        is_template_hash = (8 <= len(second_last) <= 10 and 
                           second_last.isalnum() and 
                           not any(c.isupper() for c in second_last))
        
        if is_random_suffix and is_template_hash:
            return '-'.join(parts[:-2])
    
    # 处理模板哈希长度异常的情况
    if len(parts) >= 3 and parts[-1].isalnum() and len(parts[-1]) == 5:
        return '-'.join(parts[:-2])
    
    # 基础匹配：常规的Deployment命名模式
    if len(parts) >= 2 and parts[-1].isalnum() and len(parts[-1]) >= 3:
        return '-'.join(parts[:-1])
    
    return pod_name

def parse_cpu_value(cpu_str):
    """解析 CPU 值并转换为核数"""
    if not cpu_str:
        return 0.0
    
    if isinstance(cpu_str, (int, float)):
        return float(cpu_str)
    
    if isinstance(cpu_str, str):
        if cpu_str.endswith('m'):
            return float(cpu_str[:-1]) / 1000.0
        try:
            return float(cpu_str)
        except ValueError:
            return 0.0
    
    return 0.0

def parse_memory_value(memory_str):
    """解析内存值并转换为 MB"""
    if not memory_str:
        return 0.0
    
    if isinstance(memory_str, (int, float)):
        return float(memory_str) / (1024 * 1024)
    
    if not isinstance(memory_str, str):
        return 0.0
    
    memory_str = memory_str.upper()
    if memory_str.endswith('KI'):
        return float(memory_str[:-2]) / 1024.0
    elif memory_str.endswith('MI'):
        return float(memory_str[:-2])
    elif memory_str.endswith('GI'):
        return float(memory_str[:-2]) * 1024.0
    elif memory_str.endswith('TI'):
        return float(memory_str[:-2]) * 1024.0 * 1024.0
    elif memory_str.endswith('K'):
        return float(memory_str[:-1]) / 1000.0
    elif memory_str.endswith('M'):
        return float(memory_str[:-1])
    elif memory_str.endswith('G'):
        return float(memory_str[:-1]) * 1000.0
    elif memory_str.endswith('T'):
        return float(memory_str[:-1]) * 1000.0 * 1000.0
    
    # 默认假设是字节
    try:
        return float(memory_str) / (1024 * 1024)
    except ValueError:
        return 0.0

def parse_capacity_provisioned(capacity_str):
    """从CapacityProvisioned注解中解析CPU和内存信息"""
    if not capacity_str or capacity_str == "N/A":
        return 0.0, 0.0
    
    cpu_cores = 0.0
    memory_mb = 0.0
    
    # 解析格式如: "0.25vCPU 0.5GB" 或 "1vCPU 5GB"
    import re
    
    # 匹配CPU: 数字 + vCPU
    cpu_match = re.search(r'(\d+(?:\.\d+)?)vCPU', capacity_str)
    if cpu_match:
        cpu_cores = float(cpu_match.group(1))
    
    # 匹配内存: 数字 + GB
    memory_match = re.search(r'(\d+(?:\.\d+)?)GB', capacity_str)
    if memory_match:
        memory_mb = float(memory_match.group(1)) * 1024.0  # GB转MB
    
    return cpu_cores, memory_mb

def format_cpu_value(cpu_cores):
    """格式化 CPU 值"""
    if cpu_cores == 0:
        return "0m"
    elif cpu_cores < 1:
        # 保留小数精度，避免使用int()丢失小数部分
        millis = cpu_cores * 1000
        if millis == int(millis):
            return f"{int(millis)}m"
        else:
            return f"{millis:.1f}m"
    else:
        return f"{cpu_cores:.2f} cores"

def format_memory_value(memory_mb):
    """格式化内存值"""
    if memory_mb == 0:
        return "0 Mi"
    elif memory_mb < 1024:
        # 保留小数精度，避免使用int()丢失小数部分
        if memory_mb == int(memory_mb):
            return f"{int(memory_mb)} Mi"
        else:
            return f"{memory_mb:.1f} Mi"
    else:
        return f"{memory_mb / 1024:.2f} Gi"

def main():
    # 设置 Kubernetes 客户端
    core_api, apps_api = setup_kubernetes_client()
    
    # 获取所有 Deployment 信息
    deployments = {}
    try:
        deployments_list = apps_api.list_deployment_for_all_namespaces()
        for deployment in deployments_list.items:
            key = f"{deployment.metadata.namespace}/{deployment.metadata.name}"
            deployments[key] = {
                "replicas": deployment.spec.replicas or 0,
                "available_replicas": deployment.status.available_replicas or 0
            }
    except Exception as e:
        print(f"获取 Deployment 列表失败: {e}")
    
    # 创建 Excel 工作簿
    workbook = Workbook()
    sheet = workbook.active
    sheet.title = "Deployment资源信息"
    
    # 添加表头
    headers = [
        "Namespace", "Deployment Name", "Replicas", 
        "Capacity Provisioned", "CPU Request", "CPU Limit", 
        "Memory Request", "Memory Limit"
    ]
    
    for col, header in enumerate(headers, 1):
        sheet.cell(row=1, column=col, value=header)
    
    # 获取所有 Pod
    try:
        pods = core_api.list_pod_for_all_namespaces(watch=False).items
        print(f"找到 {len(pods)} 个 Pod")
    except Exception as e:
        print(f"获取 Pod 列表失败: {e}")
        sys.exit(1)
    
    # 按 Deployment 聚合数据
    deployment_data = defaultdict(lambda: {
        "namespace": "",
        "deployment_name": "",
        "replicas": 0,
        "available_replicas": 0,
        "capacity_provisioned": set(),
        "cpu_request": 0.0,
        "cpu_limit": 0.0,
        "memory_request": 0.0,
        "memory_limit": 0.0,
        "pod_count": 0
    })
    
    # 遍历所有 Pod 并聚合信息
    for pod in pods:
        namespace = pod.metadata.namespace
        pod_name = pod.metadata.name
        
        # 提取 Deployment 名称
        deployment_name = extract_deployment_name(pod_name)
        key = f"{namespace}/{deployment_name}"
        
        # 获取注解信息
        annotations = pod.metadata.annotations or {}
        capacity_provisioned = annotations.get("CapacityProvisioned", "N/A")
        
        # 初始化资源变量
        cpu_request = cpu_limit = memory_request = memory_limit = 0.0
        
        # 遍历容器
        if pod.spec.containers:
            for container in pod.spec.containers:
                resources = container.resources or {}
                requests = resources.requests or {}
                limits = resources.limits or {}
                
                # 获取容器的资源配置
                container_cpu_request = parse_cpu_value(requests.get("cpu", "0"))
                container_cpu_limit = parse_cpu_value(limits.get("cpu", "0"))
                container_memory_request = parse_memory_value(requests.get("memory", "0"))
                container_memory_limit = parse_memory_value(limits.get("memory", "0"))
                
                # 如果容器没有设置资源配置，但Pod有CapacityProvisioned注解，尝试从注解中解析
                if (container_cpu_request == 0 and container_memory_request == 0 and 
                    container_cpu_limit == 0 and container_memory_limit == 0 and 
                    capacity_provisioned != "N/A"):
                    
                    annotation_cpu, annotation_memory = parse_capacity_provisioned(capacity_provisioned)
                    if annotation_cpu > 0 or annotation_memory > 0:
                        print(f"信息: Pod {pod_name} 从CapacityProvisioned注解解析资源: {annotation_cpu}vCPU {annotation_memory/1024:.2f}GB")
                        # 使用注解中的资源信息
                        container_cpu_request = container_cpu_limit = annotation_cpu
                        container_memory_request = container_memory_limit = annotation_memory
                
                # 累加所有容器的资源请求和限制
                cpu_request += container_cpu_request
                cpu_limit += container_cpu_limit
                memory_request += container_memory_request
                memory_limit += container_memory_limit
        
        # 更新 Deployment 数据 - 只记录单个Pod的资源配置，不累加
        data = deployment_data[key]
        data["namespace"] = namespace
        data["deployment_name"] = deployment_name
        data["capacity_provisioned"].add(capacity_provisioned)
        
        # 只记录第一个Pod的资源配置，不累加多个Pod
        if data["pod_count"] == 0:
            data["cpu_request"] = cpu_request
            data["cpu_limit"] = cpu_limit
            data["memory_request"] = memory_request
            data["memory_limit"] = memory_limit
        
        data["pod_count"] += 1
        
        # 从已知的 Deployment 信息中获取副本数
        if key in deployments:
            data["replicas"] = deployments[key]["replicas"]
            data["available_replicas"] = deployments[key]["available_replicas"]
    
    # 将数据分为业务命名空间和系统命名空间
    business_data = []
    system_data = []
    
    for key, data in deployment_data.items():
        if data["namespace"] in SYSTEM_NAMESPACES:
            system_data.append(data)
        else:
            business_data.append(data)
    
    # 对业务命名空间按名称排序
    business_data.sort(key=lambda x: (x["namespace"], x["deployment_name"]))
    
    # 对系统命名空间按指定顺序排序
    system_data.sort(key=lambda x: (
        list(SYSTEM_NAMESPACES).index(x["namespace"]) if x["namespace"] in SYSTEM_NAMESPACES else len(SYSTEM_NAMESPACES),
        x["deployment_name"]
    ))
    
    # 写入数据到 Excel
    row = 2
    
    # 先写入业务命名空间
    for data in business_data:
        # 处理 Capacity Provisioned（如果有多个值，用逗号分隔）
        capacity_provisioned = ", ".join(data["capacity_provisioned"]) if data["capacity_provisioned"] else "N/A"
        
        # 格式化副本数为 "期望/可用" 格式
        replicas_str = f"{data['replicas']}/{data['available_replicas']}"
        
        sheet.cell(row=row, column=1, value=data["namespace"])
        sheet.cell(row=row, column=2, value=data["deployment_name"])
        sheet.cell(row=row, column=3, value=replicas_str)
        sheet.cell(row=row, column=4, value=capacity_provisioned)
        sheet.cell(row=row, column=5, value=format_cpu_value(data["cpu_request"]))
        sheet.cell(row=row, column=6, value=format_cpu_value(data["cpu_limit"]))
        sheet.cell(row=row, column=7, value=format_memory_value(data["memory_request"]))
        sheet.cell(row=row, column=8, value=format_memory_value(data["memory_limit"]))
        
        row += 1
    
    # 写入系统命名空间
    for data in system_data:
        # 处理 Capacity Provisioned（如果有多个值，用逗号分隔）
        capacity_provisioned = ", ".join(data["capacity_provisioned"]) if data["capacity_provisioned"] else "N/A"
        
        # 格式化副本数为 "期望/可用" 格式
        replicas_str = f"{data['replicas']}/{data['available_replicas']}"
        
        sheet.cell(row=row, column=1, value=data["namespace"])
        sheet.cell(row=row, column=2, value=data["deployment_name"])
        sheet.cell(row=row, column=3, value=replicas_str)
        sheet.cell(row=row, column=4, value=capacity_provisioned)
        sheet.cell(row=row, column=5, value=format_cpu_value(data["cpu_request"]))
        sheet.cell(row=row, column=6, value=format_cpu_value(data["cpu_limit"]))
        sheet.cell(row=row, column=7, value=format_memory_value(data["memory_request"]))
        sheet.cell(row=row, column=8, value=format_memory_value(data["memory_limit"]))
        
        row += 1
    
    # 保存 Excel 工作簿
    try:
        workbook.save("deployments_resources.xlsx")
        print("Excel 文件已保存为 'deployments_resources.xlsx'")
    except Exception as e:
        print(f"保存 Excel 文件失败: {e}")

if __name__ == "__main__":
    # 调用主函数    
    main()